import os
import string
import secrets
import base64
import logging
import platform
from typing import Optional, Union, Tuple, List

# 配置日志
logger = logging.getLogger(__name__)

class SystemKeychain:
    """
    系统密钥管理类，负责密钥的生成、存储、检索和轮换
    """
    ENV_VAR_NAME = "SECURECONFIGMASTERKEY"
    KEY_LENGTH = 32
    
    @classmethod
    def generate_secure_key(cls) -> str:
        """
        生成加密安全的主密钥
        
        Returns:
            str: 新生成的安全密钥
        """
        alphabet = string.ascii_letters + string.digits + "!@#$%^&*()"
        key = ''.join(secrets.choice(alphabet) for _ in range(cls.KEY_LENGTH))
        return key
    
    @classmethod
    def initialize_keychain(cls) -> str:
        """
        初始化系统钥匙链，生成新密钥并存储到环境变量
        
        Returns:
            str: 新生成的主密钥
        """
        # 检查是否已初始化
        if cls.is_initialized():
            logger.warning("密钥系统已初始化，正在返回当前密钥")
            try:
                return os.environ[cls.ENV_VAR_NAME]
            except KeyError:
                logger.error(f"环境变量 {cls.ENV_VAR_NAME} 不存在")
                # 继续初始化流程
        
        # 生成新密钥
        master_key = cls.generate_secure_key()
        
        # 存储密钥到环境变量(临时)
        os.environ[cls.ENV_VAR_NAME] = master_key
        
        logger.info("密钥系统初始化成功")
        return master_key
    
    @classmethod
    def get_master_key(cls) -> bytes:
        """
        获取当前主密钥
        
        Returns:
            bytes: 当前主密钥的字节表示
        
        Raises:
            KeyError: 如果主密钥未设置
        """
        key = cls.get_key_from_environment()
        if not key:
            raise KeyError(f"未找到主密钥，请确保环境变量 {cls.ENV_VAR_NAME} 已正确设置")
        
        # 转换为字节以符合加密模块要求
        return key.encode('utf-8') if isinstance(key, str) else key
    
    @classmethod
    def get_key_from_environment(cls) -> Optional[str]:
        """
        从环境变量中获取密钥
        
        Returns:
            Optional[str]: 密钥字符串或None（如果未找到）
        """
        # 首先尝试从进程环境变量获取
        key = os.environ.get(cls.ENV_VAR_NAME)
        if key:
            return key
            
        # 如果不存在，尝试获取系统级环境变量
        try:
            if platform.system() == "Windows":
                return cls._get_key_from_windows_registry()
            else:
                # 在Linux/Mac等系统上可以尝试其他方法
                return None
        except Exception as e:
            logger.error(f"获取系统密钥失败: {str(e)}")
            return None
    
    @classmethod
    def _get_key_from_windows_registry(cls) -> Optional[str]:
        """从Windows注册表获取密钥"""
        try:
            import winreg
            
            # 尝试用户级环境变量
            try:
                key = winreg.OpenKey(
                    winreg.HKEY_CURRENT_USER,
                    'Environment',
                    0,
                    winreg.KEY_READ | winreg.KEY_WOW64_64KEY
                )
                value, _ = winreg.QueryValueEx(key, cls.ENV_VAR_NAME)
                winreg.CloseKey(key)
                return value
            except WindowsError:
                # 尝试系统级环境变量
                key = winreg.OpenKey(
                    winreg.HKEY_LOCAL_MACHINE,
                    r'SYSTEM\CurrentControlSet\Control\Session Manager\Environment',
                    0,
                    winreg.KEY_READ | winreg.KEY_WOW64_64KEY
                )
                value, _ = winreg.QueryValueEx(key, cls.ENV_VAR_NAME)
                winreg.CloseKey(key)
                return value
        except Exception as e:
            logger.error(f"从Windows注册表读取密钥失败: {str(e)}")
            return None
    
    @classmethod
    def is_initialized(cls) -> bool:
        """
        检查系统是否已初始化
        
        Returns:
            bool: 如果密钥已设置则返回True，否则返回False
        """
        return cls.get_key_from_environment() is not None
    
    @classmethod
    def rotate_key(cls, old_key: bytes) -> bytes:
        """
        执行密钥轮换操作
        
        Args:
            old_key: 当前使用的密钥
            
        Returns:
            bytes: 新生成的密钥
        """
        # 验证旧密钥
        current_key = cls.get_master_key()
        if current_key != old_key:
            logger.warning("提供的旧密钥与当前系统密钥不匹配")
        
        # 生成新密钥
        new_key_str = cls.generate_secure_key()
        new_key = new_key_str.encode('utf-8')
        
        # 更新环境变量(临时)
        os.environ[cls.ENV_VAR_NAME] = new_key_str
        
        logger.info("密钥轮换成功，请确保将新密钥添加到系统环境变量")
        return new_key
    
    @classmethod
    def set_master_key(cls, key: Union[str, bytes]) -> bool:
        """
        设置主密钥（仅用于测试或特殊情况）
        
        Args:
            key: 要设置的密钥
            
        Returns:
            bool: 设置是否成功
        """
        try:
            if isinstance(key, bytes):
                key_str = key.decode('utf-8')
            else:
                key_str = key
                
            os.environ[cls.ENV_VAR_NAME] = key_str
            return True
        except Exception as e:
            logger.error(f"设置主密钥失败: {str(e)}")
            return False
            
    @classmethod
    def clear_key(cls) -> bool:
        """
        从当前会话环境变量中清除密钥
        
        Returns:
            bool: 清除是否成功
        """
        try:
            if cls.ENV_VAR_NAME in os.environ:
                del os.environ[cls.ENV_VAR_NAME]
                logger.info("已从当前会话中清除密钥")
                return True
            else:
                logger.info("没有找到需要清除的密钥")
                return False
        except Exception as e:
            logger.error(f"清除密钥失败: {str(e)}")
            return False

    @classmethod
    def clear_user_key(cls) -> Tuple[bool, str]:
        """
        清除用户密钥但保持系统功能，生成临时密钥
        
        Returns:
            tuple: (清除是否成功, 临时密钥)
        """
        try:
            # 生成新的临时密钥
            temp_key = cls.generate_secure_key()
            
            # 清除当前密钥
            cls.clear_key()
            
            # 设置临时密钥
            os.environ[cls.ENV_VAR_NAME] = temp_key
            logger.info("已清除用户密钥并生成临时密钥")
            
            return True, temp_key
        except Exception as e:
            logger.error(f"清除用户密钥失败: {str(e)}")
            return False, ""
    
    @classmethod
    def key_exists(cls, key_id: str) -> bool:
        """
        检查指定的密钥ID是否存在
        
        Args:
            key_id: 密钥ID
            
        Returns:
            bool: 密钥是否存在
        """
        # 在这个简化版本中，我们只检查是否有主密钥
        # 在实际系统中，这里应该检查特定ID的密钥
        logger.info(f"检查密钥是否存在: {key_id}")
        return cls.is_initialized()
    
    @classmethod
    def update_key(cls, key_id: str, key_data: bytes) -> bool:
        """
        更新指定ID的密钥
        
        Args:
            key_id: 密钥ID
            key_data: 新的密钥数据
            
        Returns:
            bool: 更新是否成功
        """
        # 在这个简化版本中，我们只更新主密钥
        # 在实际系统中，这里应该更新特定ID的密钥
        try:
            logger.info(f"更新密钥: {key_id}")
            key_str = key_data.decode('utf-8') if isinstance(key_data, bytes) else key_data
            os.environ[cls.ENV_VAR_NAME] = key_str
            return True
        except Exception as e:
            logger.error(f"更新密钥失败: {str(e)}")
            return False
    
    @classmethod
    def list_keys(cls) -> List[str]:
        """
        列出所有可用的密钥ID
        
        Returns:
            List[str]: 密钥ID列表
        """
        # 在这个简化版本中，我们只返回主密钥(如果存在)
        # 在实际系统中，这里应该返回所有密钥的ID列表
        if cls.is_initialized():
            logger.info("列出所有密钥")
            return ["主密钥"]
        return []


